# - * - coding: utf - 8 - * -
#!/usr/bin/env python
import rospy
from sensor_msgs.msg import LaserScan
import rosbag
import math
import matplotlib.pyplot as plt
import numpy as np
from filterpy.kalman import KalmanFilter

previous_min_non_zero = [float('inf'), float('inf'), float('inf'), float('inf')]
tail_space_values = []  # 初始化一个列表来存储tail_space_的数据
filtered_tail_space_values = []  # 初始化一个列表来存储经过卡尔曼滤波后的tail_space_数据
kf = KalmanFilter(dim_x=1, dim_z=1)
def find_min_exclude_zero(data, status):
    min_non_zero = float('inf')
    
    for d in data:
        if d > 0 and d < min_non_zero and not math.isinf(d):
            min_non_zero = d
            
    if min_non_zero == float('inf'):
        if 1 <= status <= 4 and previous_min_non_zero[status - 1] != float('inf'):
            return previous_min_non_zero[status - 1]
        else:
            print("返回0，因为没有找到非零有限值")
            return 0
    else:
        if 1 <= status <= 4:
            previous_min_non_zero[status - 1] = min_non_zero
        return min_non_zero


# 定义一个回调函数，用于处理读取到的 LaserScan 消息
def scan_callback(laser_msg):
    global kf
    print("Ranges Length: ", len(laser_msg.ranges))
    g_ranges_ = laser_msg.ranges
    # 使用Python的切片(slice)功能代替std::min_element
    # head_full_data = g_ranges_[:180] + g_ranges_[540:]
    # head_space_ = find_min_exclude_zero(head_full_data, 1)
    # left_space_ = find_min_exclude_zero(g_ranges_[120:240], 3)
    # right_space_ = find_min_exclude_zero(g_ranges_[480:600], 4)

    tail_space_ = find_min_exclude_zero(msg.ranges[330:390], 2)
    tail_space_values.append(tail_space_)  # 将每次计算的tail_space_添加到列表中
    # 使用卡尔曼滤波器处理tail_space_数据
    print(len(tail_space_values))
    if len(tail_space_values) == 1:
        
        kf.x = np.array([tail_space_])  # 初始状态
        kf.P *= 1000  # 初始协方差矩阵
        kf.F = np.array([[1.]])  # 状态转移矩阵
        kf.H = np.array([[1.]])  # 观测矩阵
        kf.R *= 2.5 # 观测噪声协方差
        kf.Q *= 0.1  # 过程噪声协方差
    else:
        kf.predict()
        kf.update(tail_space_)
        filtered_tail_space_values.append(kf.x[0])

    # 实时绘制折线图
    plt.clf()  # 清除之前的图像
    plt.plot(tail_space_values, label='Raw Tail Space')
    plt.plot(filtered_tail_space_values, label='Filtered Tail Space')
    plt.xlabel('Time Step')
    plt.ylabel('Distance')
    plt.title('Dynamic Line Plot of Tail Space')
    plt.legend()
    plt.pause(0.001)  # 暂停一小段时间以更新图形

# 初始化 ROS 节点
rospy.init_node('scan_reader')

# 创建一个 rosbag 对象
bag = rosbag.Bag('/home/zhangjj/dm_ws/src/turn_on_robot/script/scan.bag')  # 替换为你的 rosbag 文件路径

# 创建图形窗口
plt.figure()
plt.show(block=False)  # 创建图形窗口，但不阻塞主线程

# 遍历 rosbag 中的每个消息
for topic, msg, t in bag.read_messages(topics=['/scan']):
    scan_callback(msg)  # 调用回调函数处理 LaserScan 消息
        # 检查是否按下了结束按钮
    if plt.waitforbuttonpress(0.1):  # 等待0.1秒，检查是否按下了按钮
        break  # 如果按下了按钮，退出循环

# 关闭 rosbag
bag.close()

# 阻塞主线程，直到用户关闭图形窗口
plt.ioff()
plt.show()